from mindquantum import *
import mindspore as ms
import mindspore.nn as nn
import scipy.sparse.linalg as sla
from scipy.sparse import csr_matrix
from .ham import get_ham_matrix 
import numpy as np

def hva_shallow(N, depth):
    circ = Circuit()
    circ += UN(X,[i for i in range(0, N, 2)])
    for i in range(N-1):
        circ += RZ(f'theta{i}').on(i)
    
    for d in range(depth):
        for i in range(N-1):
            if i % 2 == 0:
                circ += Rzz(f'theta{d}_{i}, {i+1}').on((i, i+1))
                circ += Rxx(f'gamma{d}_{i}, {i+1}').on((i, i+1))
                circ += Ryy(f'gamma{d}_{i}, {i+1}').on((i, i+1))
        
        for i in range(N-1):
            if i % 2 == 1:
                circ += Rzz(f'theta{d}_{i}, {i+1}').on((i, i+1))
                circ += Rxx(f'gamma{d}_{i}, {i+1}').on((i, i+1))
                circ += Ryy(f'gamma{d}_{i}, {i+1}').on((i, i+1))

    return circ 

ms.set_context(mode=ms.PYNATIVE_MODE, device_target="CPU")

# def classical_qaoa_evolution(BIG_adjacent_weight_matrix, BIG_adjacent_matrix, N, circ, epoch,depth, P=4):
#     circ = circ(N, depth)
#     sim = Simulator('mqvector', circ.n_qubits)                     # 创建模拟器，backend使用‘mqvector’，能模拟5个比特（'circ'线路中包含的比特数）
#     ham = get_ham_matrix(BIG_adjacent_weight_matrix, BIG_adjacent_matrix, P)

#     # eigenvalues, eigenvectors = sla.eigsh(H_matrix, k=1, which='SA')

#     # va, vt = np.linalg.eig(ham)
#     # min_eigenvalue_index = np.argmin(va)
#     # max_eigenvalue_index = np.argmax(va)

#     # # 获取与最小本征值对应的本征向量
#     # min_eigenvalue = va[min_eigenvalue_index]
#     # max_eigenvalue = va[max_eigenvalue_index]

#     ham = csr_matrix(ham)

#     min_eigenvalue, min_eigenvector = sla.eigsh(ham, k=1, which='SA')
#     max_eigenvalue, max_eigenvector = sla.eigsh(ham, k=1, which='LA')

#     accuracylist = []

#     grad_ops = sim.get_expectation_with_grad(Hamiltonian(ham), circ)            # 获取计算变分量子线路的期望值和梯度的算子
#     net = MQAnsatzOnlyLayer(grad_ops)                              # 生成待训练的神经网络
#     opti = nn.Adam(net.trainable_params(), learning_rate=0.05)     # 设置针对网络中所有可训练参数、学习率为0.05的Adam优化器
#     train_net = nn.TrainOneStepCell(net, opti)
#     for e in range(epoch):
#         train_net()
#         theta = net.weight.asnumpy()
#         f, _ = grad_ops(theta)
#         ratio = (max_eigenvalue-f.real)/(max_eigenvalue-min_eigenvalue)
#         accuracylist.append(ratio[0][0])
    
#     return theta, accuracylist, circ

def get_coefficient_matrix(BIG_vertex_weight_matrix, BIG_adjacent_matrix, P=4):
    nqubits = len(BIG_adjacent_matrix)

    J = np.zeros((nqubits, nqubits))
    h = np.zeros((nqubits))

    w_list = BIG_vertex_weight_matrix.flatten()
    
    for i in range(nqubits):
        for j in range(nqubits):
            if BIG_adjacent_matrix[i, j] == 0 and i!=j:
                J[i, j] += P
                h[i] -= P/2
                h[j] -= P/2

    
    for i in range(nqubits):
        h[i] += w_list[i]/2
    
    h = h[:, np.newaxis]
    return J, h

# QAOA definition
def qaoa(N, depth):
    circ = Circuit()
    circ += UN(H, N) # prepare initial state

    for p in range(depth):
        # generate U_C
        for i in range(N-3):
            circ += Rzz(f'rzz_{p}{i}').on((i, i+3))
        for i in range(N):
            circ += RZ(f'rz_{p}{i}').on(i)

        # generate U_B
        for i in range(N):
            circ += RX(f'beta_{p}{i}').on(i)
    return circ

def dc_qaoa(N, depth):
    circ = Circuit()
    circ += UN(H, N) # prepare initial state

    for p in range(depth):
        # generate U_C
        for i in range(N-3):
            circ += Rzz(f'rzz_{p}{i}').on((i, i+3))
        for i in range(N):
            circ += RZ(f'rz_{p}{i}').on(i)
        # generate U_B
        for i in range(N):
            circ += RX(f'rx_{p}{i}').on(i)
        # generate CD part
        for i in range(N):
            circ += RY(f'ry_{p}{i}').on(i)
            
    return circ

def abf_qaoa(N, depth):
    circ = Circuit()
    # circ += UN(H, N) # prepare initial state

    for i in range(N):
        circ += RZ(f'init_bias_{i}').on(i)

    for p in range(depth):
        # generate U_C
        for i in range(N-3):
            circ += Rzz(f'rzz_{p}{i}').on((i, i+3))
        for i in range(N):
            circ += RZ(f'rz_{p}{i}').on(i)
        # generate U_B
        for i in range(N):
            circ += RX(f'rx_{p}{i}').on(i)
        # generate CD part
        for i in range(N):
            circ += RY(f'ry_{p}{i}').on(i)
            
    return circ